import os
from logger import ULogger
import queue
import json
from apscheduler.schedulers.blocking import BlockingScheduler
import threading
from multiprocessing import Queue,Process,Manager
import time
import consul

def get_subfolders(folder):
    for _, dirs, _ in os.walk(folder):
       return dirs

class ModelRepoMonitor:
    def __init__(self,repo,consul_host=None,consul_port=8500):
        self.repo=repo
        self.task_cache=Manager().dict()
        self.task_done_time_cache = Manager().dict()
        self.task_queue=Manager().Queue()
        self.logger=ULogger("log.txt")
        self.running=True
        self.interval_time=30
        self.consul=None
        if consul_host is not None:
            try:
                self.consul = consul.Consul(consul_host, consul_port)
            except Exception as e:
                print(e)

    def process_task(self):
        while self.running:
            try:
                model_name=None
                task = self.task_queue.get(timeout=60*5)
                if task is None:
                    continue
                model_name, version, backend, func = task
                config_path=os.path.join(self.repo,model_name,version,backend,"convert_config.json")
                ret = func(config_path)
                if ret!=0:
                    self.logger.error("task failed {0} {1} {2}".format(model_name,version,backend))
                self.task_done_time_cache[(model_name,version,backend)]=int(time.time())
                self.task_cache.pop((model_name,version,backend))
            except Exception as e:
                if model_name is not None:
                    self.task_cache.pop((model_name, version, backend))
                self.logger.error("task failed {0}".format(e))
                continue

    def register_rknn_quant_func(self,rknn_quant_func):
        self.rknn_quant=rknn_quant_func

    def monitor_model(self, model_name):
        versions= get_subfolders(os.path.join(self.repo, model_name))
        for version in versions:
            backends = get_subfolders(os.path.join(self.repo, model_name, version))
            if 'rknn' in backends and (model_name,version,'rknn') not in self.task_cache:
                #check if we need to quant
                config_path = os.path.join(self.repo, model_name, version, 'rknn', "convert_config.json")
                if not os.path.exists(config_path):
                    continue

                #check if rknn model already exist
                if os.path.exists(os.path.join(self.repo,model_name,version,'rknn','{0}.rknn'.format(version))):
                    continue
                #check if we have quant this model recently
                if (model_name,version,'rknn') in self.task_done_time_cache:
                    current_time=int(time.time())
                    if current_time-self.task_done_time_cache[((model_name,version,'rknn'))] < self.interval_time:
                        continue
                    self.task_done_time_cache.pop((model_name,version,"rknn"))

                self.logger.info("add rknn task in queue: {0} {1}".format(model_name,version))
                self.task_cache[(model_name,version,'rknn')]=1
                self.task_queue.put((model_name,version,'rknn',self.rknn_quant))

    def monitor_repo(self):
        def monitor_once():
            dirs=get_subfolders(self.repo)
            for model in dirs:
                self.monitor_model(model)
        def set_interval_time():
            try:
                value = self.consul.kv.get("settings/interval")
                interval = int(value[1]['Value'])*2
                if interval > 0 and interval != self.interval:
                    self.interval = interval
                    self.sync_job.reschedule("interval", seconds=interval)
                    print("change interval to {0}".format(interval))
            except Exception as e:
                pass
        self.sched = BlockingScheduler()
        self.sync_job = self.sched.add_job(monitor_once, 'interval', args=(), id='sync_repo', seconds=self.interval_time)
        if self.consul is not None:
            self.sched.add_job(set_interval_time, 'interval', args=(), id='set_interval', seconds=self.interval_time)
        self.sched.start()

    def start_monitor(self):
        self.process_thread=Process(target=self.process_task,args=())
        self.monitor_thread = threading.Thread(target=self.monitor_repo, args=())
        self.monitor_thread.start()
        self.process_thread.start()
        self.process_thread.join()

    def stop_monitor(self):
        self.sched.remove_all_jobs()
        self.sched.shutdown()
        self.running=False

    def __del__(self):
        try:
            self.stop_monitor()
        except Exception:
            pass

